package com.example.Arithmetic.Arithmetic;

import java.util.Comparator;
import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;

/**
 * 日期：2023/12/31
 * 时间：9:03
 * 描述：哈夫曼树，用来编码解码，是个满二叉树
 */
public class HuffmanTree {
    class Node {
        Character ch;    //字符
        Node left;
        Node right;
        int freq;      //频次
        String code;   //编码

        public Node(Node left, Node fight, int freq) {
            this.left = left;
            this.right = fight;
            this.freq = freq;
        }

        public Node(Character ch) {
            this.ch = ch;
        }

        int freq() {
            return freq;
        }

        public boolean isLeft() {
            return left == null;
        }

        @Override
        public String toString() {
            return "Node{" +
                    "ch=" + ch +
                    ", freq=" + freq +
                    '}';
        }
    }

    String str;
    Map<Character, Node> map = new HashMap();
    Node root;

    public HuffmanTree(String str) {
        this.str = str;
        for (char c : str.toCharArray()) {
//            if (!map.containsKey(c)) {
//                map.put(c, new Node(c));
//            }
//            Node node = map.get(c);
//            node.freq++;
//            可以简化为下面
            Node node = map.computeIfAbsent(c, Node::new);
            node.freq++;
        }
//        构造一个树
        PriorityQueue<Node> priorityQueue = new PriorityQueue<>(
                Comparator.comparingInt(Node::freq));
        priorityQueue.addAll(map.values());
        while (priorityQueue.size() >= 2) {
            Node x = priorityQueue.poll();
            Node y = priorityQueue.poll();
            priorityQueue.offer(new Node(x, y, x.freq + y.freq));
        }
        root = priorityQueue.poll();
//       计算每个字符的编码
        int sum = dfs(root, new StringBuilder());
        for (Node node : map.values()) {
            System.out.println(node + "" + node.code);
        }
//        计算总共占有多少比特
        System.out.println("总共占有多少字节：" + sum);
    }

    private int dfs(Node node, StringBuilder stringBuilder) {
        int sum = 0;
        if (node.isLeft()) {
//            找到编码，给对象节点赋值
            node.code = stringBuilder.toString();
            sum = node.freq * stringBuilder.length();
        } else {
            sum += dfs(node.left, stringBuilder.append("0"));
            sum += dfs(node.right, stringBuilder.append("1"));
        }
        if (stringBuilder.length() > 0) {
            stringBuilder.deleteCharAt(stringBuilder.length() - 1);
        }
        return sum;
    }

    //    编码
    public String encode() {
        char[] chars = str.toCharArray();
        StringBuilder sb = new StringBuilder();
        for (char c : chars) {
            sb.append(map.get(c).code);
        }
        return sb.toString();
    }

    // 解码
    public String decode(String str) {
          /*
            从根节点，寻找数字对应的字符
                数字是 0 向左走
                数字是 1 向右走
                如果没走到头，每走一步数字的索引 i++
            走到头就可以找到解码字符，再将 node 重置为根节点
         */
        StringBuilder sb = new StringBuilder();
        Node node = root;
        for (char c : str.toCharArray()) {
            if (!node.isLeft()) {
                if (c == '0') {
                    node = node.left;
                } else if (c == '1') {
                    node = node.right;
                }
            }
            if (node.isLeft()) {
                sb.append(node.ch);
                node = root;
            }
        }
        return sb.toString();
    }

    public static void main(String[] args) {
        HuffmanTree huffmanTree = new HuffmanTree("abcdefghjklmnopqrstuvwsyz");
        System.out.println(huffmanTree.encode());
        String decode = huffmanTree.decode("111100010111000001");
        System.out.println(decode);

        String encode = huffmanTree.encode();
        System.out.println(huffmanTree.decode(encode));
    }
}
